from pathlib import Path
import re


class DirNotExist(Exception):
    pass


def get_python_file(directory):
    path = directory if isinstance(directory, Path) else Path(directory)

    if not path.exists():
        raise DirNotExist

    for each_path in path.iterdir():
        if each_path.is_dir():
            get_python_file(each_path)

        if '.py' in each_path.name:
            yield each_path.name


class PythonFile:
    def __init__(self, directory):
        self.__path = Path(directory)
        self.__python_files = []

    def _get_python_files(self, path):
        for each_path in path.iterdir():
            if each_path.is_dir():
                self._get_python_files(each_path)

            if '.py' in each_path.name:
                self.__python_files.append(each_path.name)

    def __iter__(self):
        self._get_python_files(self.__path)
        yield from self.__python_files


if __name__ == '__main__':
    p = PythonFile(".")
    for f in p:
        print(f)

# 判断是否是.py文件，可以试下endswith方法
